import numpy as np
import math
# 这个文件是和算法对接的 转换成指令用的
# 前面说的函数就是这里用的
from agent.agent import Agent
from env.env_cmd import EnvCmd, Point
from env.env_def import UnitType, UnitStatus, RED_AIRPORT_ID
from agent.red.RL_brain import TD3

'''
若目标处于己方探测范围的60%-100%之内，则只能发现目标，情报信息里该目标属性为不明；若目标处于己方探测范围的60%之内，则可准确识别出目标。
再qb中需要对未知飞行物进行判别。
'''

blue_num = 1 # 蓝方飞机数目
red_num = 5 # 红方飞机数目
blue_env_size = 2 # 蓝方环境中的数据，这里是x，y坐标两个数据
red_env_size =2 # 红方环境中的数据，这里是x，y坐标两个数据

s_dim = red_env_size*red_num+blue_env_size*blue_num  # 状态个数，用于RL训练
a_dim = red_num  # 动作个数（一个飞机一个动作，表示飞行方向）

a_bound = math.pi  # 动作的最高值，乘以这个范围，由于方向范围为-pi到pi，神经网络输出-1到1，那么就乘以pi
delay = 10
delay_in_action = 1
class RedAgent(Agent):
    def __init__(self, name, config, **kwargs):
        super().__init__(name, config['side'])

        self._init()

    def _init(self):
        self.RL = TD3(a_dim, s_dim, a_bound,delay =delay)
        self.total_steps = 0
        self.steps = 0
        self.total_reward = 0
        self.last_my_a2a = 0
        self.last_en_a2a = 0
        self.last_obs = np.zeros((1, s_dim), dtype=np.float32)

        self.dealy_action_arr = np.zeros((delay_in_action, a_dim), dtype=np.float32)

        # 当前状态
        self.S_list = np.zeros((delay*2+1, s_dim), dtype=np.float32)
        # 都是之前的数据，之前的状态，动作，奖励
        self.S_last_list = np.zeros((delay*2+1, s_dim), dtype=np.float32)
        self.action_list = np.zeros((delay*2+1, a_dim), dtype=np.float32)
        self.reward_list = np.zeros((delay*2+1, 1), dtype=np.float32)
        self.ar_list_pointer = 0
        self.last_action = np.zeros((a_dim), dtype=np.int)
        self.version = 1
        self.start_id = 0
        self.final_dis = 0
        self.find = False
        if self.start_id:
             print("载入模型")
             self.RL.load_model(self.version, self.start_id)

    def reset(self):
        self.steps = 0
        self.total_reward = 0
        self.last_my_a2a = 0
        self.last_en_a2a = 0
        self.S_list = np.zeros((delay*2+1, s_dim), dtype=np.float32)
        self.S_last_list = np.zeros((delay*2+1, s_dim), dtype=np.float32)
        self.action_list = np.zeros((delay*2+1, a_dim), dtype=np.float32)
        self.reward_list = np.zeros((delay*2+1, 1), dtype=np.float32)
        self.ar_list_pointer = 0
        self.last_obs = np.zeros((1, s_dim), dtype=np.float32)
        self.final_dis = 0
        self.find = False

    def step(self, sim_time,num_lun, obs, **kwargs):
        cmd_list = []
        obs_red = obs["red"]

        '''
        20001自己的飞机场，30001敌人的飞机场，这个情报是从始至终可以获得的
        如果除了机场之外有其他情报，则表示发现了飞机
        '''
        for unit in obs_red['qb']:
            if(unit['ID']!=20001 and unit['ID']!=30001):
                self.find = True
                break

        '''
        AWCS是侦察机，这里使用侦察机的原因是测试时候发现侦察机比歼灭机侦察范围大且稳定
        AIR是歼灭机
        '''
        if obs_red["airports"][0]["AWCS"] > 0 and sim_time > 0:
            # 机场控制飞机起飞
            '''
            return [EnvCmd.make_takeoff_areapatrol(RED_AIRPORT_ID, min(1, obs_red["airports"][0]["AIR"]),
                    UnitType.A2A, 100000, 0, 8000, 0, 20000, 20000, 220, 7200)]
            [x y z]
            区域长轴与正北方向角度、长度和宽度
            220为巡逻速度（浮点型，单位：米/秒,歼击机参考速度250m/s,预警机/干扰机/轰炸机参考速度166m/s,无人机参考速度90m/s）
            '''
            return [EnvCmd.make_takeoff_areapatrol(RED_AIRPORT_ID, min(1, obs_red["airports"][0]["AWCS"]), \
                                                   UnitType.AWACS, 50000, -100000+10000*obs_red["airports"][0]["AWCS"], 8000, 0, 20000, 20000, 1000, 7200)]
        '''
        打开侦察机的雷达（默认应该是打开的，防止没有打开的情况发生）
        '''
        for unit in obs_red["units"]:
            if (unit["LX"] == UnitType.AWACS):
                cmd_list.append(
                        EnvCmd.make_awcs_radarcontrol(unit['ID'],1))

        if self.find:
            input_obs = np.zeros((1, s_dim), dtype=np.float32)
            self_a2a = []
            enemy_a2a = []
            i = 0
            my_a2a = 0
            en_a2a = 0
            # 遍历红方的units，统计己方飞机数目
            for unit in obs_red["units"]:
                if (unit["LX"] == UnitType.AWACS):
                    input_obs[0][i] = unit["X"] / 100000
                    input_obs[0][i+1] = unit["Y"] / 100000
                    i += red_env_size
                    my_a2a += 1
                    self_a2a.append(unit)
            i = red_env_size*red_num # 如果有飞机阵亡，则跳过空域部分，这样保证i是指向敌方飞机的数据
            # 遍历红方的units，统计敌方飞机数目（18指向未知对象）
            for unit in obs_red["qb"]:
                if (unit["LX"] == UnitType.AWACS or unit["LX"] == 18):
                    input_obs[0][i] = unit["X"] / 100000
                    input_obs[0][i+1] = unit["Y"] / 100000
                    i += blue_env_size
                    enemy_a2a.append(unit)
                    en_a2a += 1

            if num_lun > 10:
                action = self.RL.choose_action(input_obs, self.S_list, self.S_last_list, self.action_list,
                                               self.reward_list)  # 其实这里应该要变为动作序列，先不动
            else:
                # 看到就追
                # action = self.get_theta(enemy_a2a[0], self_a2a[0])
                action, trace_id, _ = self.My_way(self_a2a, enemy_a2a, 5, 1000)

            self.action_push(action)
            action = self.dealy_action_arr[0] # 动作延迟
            print("动作：", self.dealy_action_arr)
            print("选择动作：", action)

            # 定义一个direction_sum，通过这个的值来判断奖励函数
            # 如果飞机一直往右跑，direction为0，1，7这几个值，那么扣分
            # 往左跑，direction为3，4，5，加分
            # direction_sum = 0

            reward_dis = 0
            reward_dir = 0
            distance = 0
            for i in range(my_a2a):
                i_unit = self_a2a[i]
                input_obs[0][i * 2]
                input_obs[0][i * 2 + 1]
                direction = action[i]
                # 此处direction是-pi到pi的

                print("direction: ", direction)
                # print("x, y: ", math.cos(direction), math.sin(direction))
                cmd_list.append(
                    EnvCmd.make_areapatrol(i_unit["ID"], self_a2a[i]["X"] + math.cos(direction)*20000,
                                           self_a2a[i]["Y"] + math.sin(direction)*20000, i_unit["Z"], 0,
                                           1000, 1000, 1000, 7200))
                np.random.shuffle(enemy_a2a)
                for j_unit in enemy_a2a:
                    # 下面用距离来判断奖励值
                    distance = self.get_distance(i_unit, j_unit)

                    self.final_dis = distance
                    if distance < 20000:
                        # reward_dis += 0.1
                        reward_dis += (20000 - distance) * 0.00007
                    # elif distance < 50000:
                    #     reward_dis += 0.3
                    elif distance > 20000 and sim_time > 1500:
                        reward_dis -= 2

                    # 下面用飞机前进的角度来判定奖励值
                    # 计算以红方为原点，做射线射向蓝方，该射线的角度值（弧度制），向左为-pi或pi，向右为0
                    # 这是理想情况下，红方应该去追击的方向
                    if sim_time > 1200:
                        dir_red = math.atan2(j_unit["Y"] - i_unit["Y"], j_unit["X"] - i_unit["X"])
                        reward_dir = math.pi - math.fabs(direction - dir_red)

            reward = 0
            # print("last_my_a2a", self.last_my_a2a)
            if self.steps != 0:
                # 计算延时奖励
                # 定义向前走的奖励值，越往前走奖励值越高
                # reward = 0.01 * (2400 - sim_time) + reward_dis
                reward = reward_dis + reward_dir  # 奖励值只和距离和时间有关
                if reward != 0:
                    print("reward", reward)
                self.RL.store_transition(self.last_obs[0], self.last_action / math.pi, reward, input_obs[0])
                # print("temp_obs", temp_obs)
                # print("temp_last_obs", temp_last_obs)
                # if self.total_steps % 10 == 0:
                print("total_steps: ", self.total_steps)
            if self.total_steps % 200 == 0 and self.total_steps:
                print("save model")
                self.RL.save_model(self.version, self.start_id + self.total_steps // 200)

            # if self.total_steps > 5:
            if self.total_steps > 100:
                # print("开始训练")
                self.RL.learn(num_lun,self.steps)

            self.last_my_a2a = my_a2a
            self.last_en_a2a = en_a2a
            self.store_action_reward(self.last_action/ math.pi,reward/ 10,self.last_obs[0],input_obs[0])

            self.last_obs = input_obs # 应该要变得
            self.last_action = action
            self.total_steps += 1
            self.steps += 1
            self.total_reward += reward
            self.find = False

        return cmd_list

    def get_average_reward(self):
        return self.total_reward / self.steps

    def get_distance(self, unit1, unit2):
        d_x = unit1["X"] - unit2["X"]
        d_y = unit1["Y"] - unit2["Y"]
        return math.sqrt(d_x ** 2 + d_y ** 2)

    def get_theta(self, unit1, unit2):
        d_x = unit1["X"] - unit2["X"]
        d_y = unit1["Y"] - unit2["Y"]
        return math.atan2(d_y, d_x)

    def get_theta_list(self, unit1, unit2):
        unit2 = list(unit2)
        d_x = unit1[0] - unit2[0]
        d_y = unit1[1] - unit2[1]
        return math.atan2(d_y, d_x)
    def get_final_dis(self):
        return self.final_dis

    def store_action_reward(self,a, r,s,s_):
        # transition = np.hstack((s, a_list, r_list, s_))
        if self.ar_list_pointer<delay*2+1:
            self.action_list[self.ar_list_pointer, :] = a
            self.reward_list[self.ar_list_pointer, :] = r
            self.S_last_list[self.ar_list_pointer - 1, :] = s
            self.S_list[self.ar_list_pointer, :] = s_
            self.ar_list_pointer = self.ar_list_pointer+1
        else:
            # [...-2，-1，最新]
            self.action_list[0:-1, :] = self.action_list[1:, :]
            self.action_list[-1,:] = a
            self.reward_list[0:-1, :] = self.reward_list[1:, :]
            self.reward_list[-1,:] = r
            self.S_last_list[0:-1, :] = self.S_last_list[1:, :] # 最后一个是未知的
            self.S_last_list[-1, :] = s
            self.S_list[0:-1, :] = self.S_list[1:, :]
            self.S_list[-1, :] = s_
    def action_push(self,a):
        # [...-2，-1，最新]
        self.dealy_action_arr[0:-1, :] = self.dealy_action_arr[1:, :]
        self.dealy_action_arr[-1, :] = a

    # 计算两点之间的距离
    def calculate_distance(self, point1, point2):
        '''
        函数说明:
        计算两个坐标点之间的距离。

        参数说明：
        point1、point2：
            目标点:[x,y]

        返回值：
            两个坐标之间的欧氏距离
        '''
        return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)

    def surrounding_points(self, enemy_drones_xy, num_circle, surrounding_radius):
        '''
        包围的目标是围在目标飞机的四周，可以将目标位置设置为四周的各个点位，让飞机前往距离最近的点位，只要让所有飞机飞行距离最短就可以了。
        输入：敌方飞机的坐标list：[[],[],...]  每个飞机包围目标个数：int
        输出：敌方飞机周围包围点list:[[],[],...]
        '''
        num_enemy_drones = len(enemy_drones_xy)
        angle_step = 2 * math.pi / num_circle
        points_arrive = []
        for i in range(num_enemy_drones):
            for j in range(num_circle):
                angle = j * angle_step
                x = enemy_drones_xy[i][0] + surrounding_radius * math.cos(angle)
                y = enemy_drones_xy[i][1] + surrounding_radius * math.sin(angle)
                # 判断坐标是否为0
                if abs(x) < 0.000000001:
                    x = 0.0
                if abs(y) < 0.000000001:
                    y = 0.0
                points_arrive.append([x, y])
        return points_arrive

    def find_nearest_drone(self, target_point, self_drones):
        '''
        函数说明:
        计算飞机集群到目标点最短距离的飞机下标。

        参数说明：
        target_point：
            目标点:[x,y]
        self_drones ：
            可用的飞机集合，最初来自于 obs[red][units]，其中数据需要使用['X']等方法得到，并不是一个现成的列表list

        返回值：
        drones_index[0]:
            己方飞机集合中最近的飞机下标 int
        drones_xy[0]
            己方飞机集合中最近的飞机坐标 list[x,y]
        drones_id[0]
            己方飞机集合中最近的飞机ID int
        '''
        drones_with_indices = [(index, (drone['X'], drone['Y']), drone['ID']) for index, drone in
                               enumerate(self_drones)]
        if not drones_with_indices:
            return []  # 返回一个空列表或适当的默认值，以防止空列表引发错误
        sorted_drones = sorted(drones_with_indices,
                               key=lambda item: self.calculate_distance(target_point, item[1]))  # 从小到大对目标点距离进行排序
        drones_index = [item[0] for item in sorted_drones]
        drones_xy = [item[1] for item in sorted_drones]
        drones_id = [item[2] for item in sorted_drones]

        return drones_index[0], drones_xy[0], drones_id[0]

    def get_surrounding_points(self, enemy_drones_xy, num_circle, surrounding_radius):
        '''
        函数说明:
        包围的目标是围在目标飞机的四周，可以将目标位置设置为四周的各个点位，让飞机前往距离最近的点位，只要让所有飞机飞行距离最短就可以了。
        前提是，己方飞机大于等于敌方飞机*包围需要的飞机
        参数说明：
        enemy_drones_xy
            敌方飞机的坐标 list：[[x,y],[x,y],...]
        num_circle
            每个敌方飞机周围包围的飞机数目 int
        surrounding_radius
            包围半径 float

        返回值：
        points_arrive
            敌方飞机周围包围点list:[[x,y],[x,y],...]

        eg:
            surrounding_points([[0,0],[1,-2]],1, 10)
            [[10.0, 0.0], [11.0, -2.0]]
        '''
        num_enemy_drones = len(enemy_drones_xy)
        angle_step = 2 * math.pi / num_circle
        points_arrive = []
        for i in range(num_enemy_drones):
            for j in range(num_circle):
                angle = j * angle_step
                x = enemy_drones_xy[i][0] + surrounding_radius * math.cos(angle)
                y = enemy_drones_xy[i][1] + surrounding_radius * math.sin(angle)
                # 判断坐标是否为0
                if abs(x) < 0.000000001:
                    x = 0.0
                if abs(y) < 0.000000001:
                    y = 0.0
                points_arrive.append([x, y])
        return points_arrive

    def calculate_surrounding_path(self, points_arrive, self_drones, circle_type):
        '''
        函数说明:
        从get_surrounding_points得到需要到达的目标后，此函数用于分配每一己方飞机需要前往的地点。考虑不同的包围方式，提供接口。

        参数说明：
        points_arrive：
            敌方飞机周围包围点list:[[],[],...]
        self_drones ：
            可用的飞机集合，最初来自于 obs[red][units]，其中数据需要使用['X']等方法得到，并不是一个现成的列表list
        circle_type ：
            str类型数据“SHORT_PATH_SELF”、“SHORT_PATH_ALL”
            "SHORT_PATH_SELF"表示己方飞单个飞机前进路程最短，如果已经有飞机前往则进入第二近。实现简单，但是容易导致并不是最近的飞机前往，并非最优解。
            "SHORT_PATH_ALL"表示己方所有飞机前进的路程最短。实现复杂，每个飞机并非前往最近点，但理论上是较好的解。

        返回值：
        Action
            己方飞机动作列表list：[[f],[f],[f]] f表示浮点数，按照正常的角坐标系
        trace_index
            己方飞机的序列
        '''
        Action = np.zeros(red_num)
        trace_index = np.zeros(red_num)
        trace_point = []
        if circle_type == "SHORT_PATH_SELF":
            units = self_drones
            for index_ene, enemy_point_circle in enumerate(points_arrive):
                if not units:
                    return Action, trace_index, trace_point
                self_index, self_point, self_ID = self.find_nearest_drone(enemy_point_circle, units)
                for index_sel, drone in enumerate(self_drones):
                    if drone['ID'] == self_ID and Action[index_sel] == 0:
                        Action[index_sel] = self.get_theta_list(enemy_point_circle, self_point)
                        trace_index[index_sel] = index_ene
                        trace_point.append(enemy_point_circle)
                        units = [units[i] for i in range(len(units)) if i != self_index]  # 逐渐减小可以使用的己方飞机数目

        return Action, trace_index, trace_point

    '''
    self_a2a和enemy_a2a的格式都是[[x,y],[x,y],...]
    '''
    def My_way(self, self_a2a, enemy_a2a, num_circle, r):
        enemy_drones_xy = []
        for index, drone in enumerate(enemy_a2a):
            enemy_drones_xy.append([drone['X'], drone['Y']])

        points_arrive = self.surrounding_points(enemy_drones_xy, num_circle, r)
        Action, trace_id, trace_point = self.calculate_surrounding_path(points_arrive, self_a2a, "SHORT_PATH_SELF")
        return Action, trace_id, trace_point